#ifndef __TWCLML_DRV_H__
#define __TWCLML_DRV_H__

#include <linux/types.h>
#include <linux/ioctl.h>

/* 定义布尔类型 */
#ifndef bool
#define bool int
#define true 1
#define false 0
#endif

/* PCI信息结构体 */
typedef struct {
    char busId[16];                /* PCI总线ID，格式如"0000:01:00.0" */
    __u32 domain;                  /* PCI域 */
    __u32 bus;                     /* PCI总线号 */
    __u32 device;                  /* PCI设备号 */
    __u32 function;                /* PCI功能号 */
} nvmlPciInfo_t;

/* 启用状态枚举 */
typedef enum {
    NVML_DISABLED = 0,             /* 已禁用 */
    NVML_ENABLED = 1               /* 已启用 */
} nvmlEnableState_t;

/* P2P能力索引枚举 */
typedef enum {
    NVML_P2P_CAPS_INDEX_READ = 0,      /* P2P读能力 */
    NVML_P2P_CAPS_INDEX_WRITE,         /* P2P写能力 */
    NVML_P2P_CAPS_INDEX_NVLINK,        /* NVLink上的P2P能力 */
    NVML_P2P_CAPS_INDEX_ATOMICS,       /* P2P原子操作能力 */
    NVML_P2P_CAPS_INDEX_PROP,          /* P2P属性 */
    NVML_P2P_CAPS_INDEX_UNKNOWN        /* 未知能力 */
} nvmlGpuP2PCapsIndex_t;

/* P2P状态枚举 */
typedef enum {
    NVML_P2P_STATUS_OK = 0,                        /* 支持且启用 */
    NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED,          /* 芯片组不支持 */
    NVML_P2P_STATUS_GPU_NOT_SUPPORTED,             /* GPU不支持 */
    NVML_P2P_STATUS_NOT_SUPPORTED,                 /* 不支持 */
    NVML_P2P_STATUS_UNKNOWN                        /* 未知状态 */
} nvmlGpuP2PStatus_t;

/* NVLink能力枚举 */
typedef enum {
    NVML_NVLINK_CAP_P2P_SUPPORTED = 0,     /* 支持NVLink上的P2P */
    NVML_NVLINK_CAP_SYSMEM_ACCESS = 1,     /* 支持访问系统内存 */
    NVML_NVLINK_CAP_P2P_ATOMICS = 2,       /* 支持P2P原子操作 */
    NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3,    /* 支持系统内存原子操作 */
    NVML_NVLINK_CAP_SLI_BRIDGE = 4,        /* 支持SLI桥接 */
    NVML_NVLINK_CAP_VALID = 5,             /* 链路有效 */
    NVML_NVLINK_CAP_COUNT                  /* 能力总数 */
} nvmlNvLinkCapability_t;

/* 错误码枚举 */
typedef enum {
    NVML_SUCCESS = 0,              /* 成功 */
    NVML_ERROR_UNINITIALIZED = 1,  /* 未初始化 */
    NVML_ERROR_INVALID_ARGUMENT = 2, /* 参数无效 */
    NVML_ERROR_NOT_SUPPORTED = 3,  /* 不支持 */
    NVML_ERROR_NOT_FOUND = 4,      /* 未找到 */
    NVML_ERROR_UNKNOWN = 999       /* 未知错误 */
} nvmlReturn_t;

/* 远程PCIe信息结构体 */
typedef struct {
    unsigned int numLinks;              /* 远程连接数 */
    nvmlPciInfo_t remotePci[4];         /* 远程PCI信息数组 */
    nvmlEnableState_t linkState[4];     /* 远程链路状态 */
    struct {
        unsigned int p2pSupported;      /* P2P支持 */
        unsigned int sysmemAccess;      /* 系统内存访问支持 */
        unsigned int p2pAtomics;        /* P2P原子操作支持 */
        unsigned int sysmemAtomics;     /* 系统内存原子操作支持 */
        unsigned int sliBridge;         /* SLI桥接支持 */
        unsigned int valid;             /* 链路有效性 */
    } linkCaps[4];                      /* 链路能力数组 */
    struct {
        nvmlGpuP2PStatus_t read;        /* 读状态 */
        nvmlGpuP2PStatus_t write;       /* 写状态 */
        nvmlGpuP2PStatus_t nvlink;      /* NVLink状态 */
        nvmlGpuP2PStatus_t atomics;     /* 原子操作状态 */
        nvmlGpuP2PStatus_t prop;        /* 属性状态 */
    } p2pCaps[4];                       /* P2P能力数组 */
} twcl_remotePcieInfo_t;

/* 设备结构体 */
typedef struct twcl_nvmlDevice {
    unsigned int index;                 /* 设备索引 */
    nvmlPciInfo_t pciInfo;              /* 本地PCI信息 */
    int computeMajor;                   /* 计算能力主版本 */
    int computeMinor;                   /* 计算能力次版本 */
    unsigned int numNvLinks;            /* NVLink数量 */
    bool chipsetSupportsP2P;            /* 芯片组是否支持P2P */
    bool gpuSupportsP2P;                /* GPU是否支持P2P */
    bool iohTopologySupportsP2P;        /* IOH拓扑是否支持P2P */
    bool p2pDisabledByRegKey;           /* 是否被注册表禁用P2P */
    twcl_remotePcieInfo_t remotePcie;   /* 远程PCIe信息（新增成员） */
} *nvmlDevice_t;

/* 设备名称和节点路径 */
#define TWCLML_DEVICE_NAME "twclml"
#define TWCLML_DEVICE_PATH "/dev/twclml"

/* IOCTL命令幻数 */
#define TWCLML_IOCTL_MAGIC 't'

/* 设备信息结构体（用户空间交互） */
typedef struct {
    unsigned int index;              /* 设备索引 */
    char bus_id[16];                 /* PCI总线ID */
    __u32 domain;                    /* PCI域 */
    __u32 bus;                       /* PCI总线号 */
    __u32 device;                    /* PCI设备号 */
    __u32 function;                  /* PCI功能号 */
    int compute_major;               /* 计算能力主版本 */
    int compute_minor;               /* 计算能力次版本 */
    unsigned int num_nvlinks;        /* NVLink数量 */
} twclml_device_info_t;

/* NVLink链路信息结构体（用户空间交互） */
typedef struct {
    unsigned int device_index;       /* 设备索引 */
    unsigned int link;               /* 链路编号 */
    int is_active;                   /* 链路是否激活 (1:激活, 0:未激活) */
    char remote_bus_id[16];          /* 远程设备PCI总线ID */
    __u32 remote_domain;             /* 远程设备PCI域 */
    __u32 remote_bus;                /* 远程设备PCI总线号 */
    __u32 remote_device;             /* 远程设备PCI设备号 */
    __u32 remote_function;           /* 远程设备PCI功能号 */
} twclml_nvlink_info_t;

/* P2P能力信息结构体（用户空间交互） */
typedef struct {
    unsigned int device1_index;      /* 第一个设备索引 */
    unsigned int device2_index;      /* 第二个设备索引 */
    int read_supported;              /* P2P读能力 */
    int write_supported;             /* P2P写能力 */
    int nvlink_supported;            /* NVLink上的P2P能力 */
    int atomics_supported;           /* P2P原子操作能力 */
} twclml_p2p_info_t;

/* IOCTL命令定义 */
#define TWCLML_IOCTL_GET_DEVICE_COUNT _IOR(TWCLML_IOCTL_MAGIC, 0, unsigned int)
#define TWCLML_IOCTL_GET_DEVICE_INFO  _IOWR(TWCLML_IOCTL_MAGIC, 1, twclml_device_info_t)
#define TWCLML_IOCTL_GET_NVLINK_INFO  _IOWR(TWCLML_IOCTL_MAGIC, 2, twclml_nvlink_info_t)
#define TWCLML_IOCTL_GET_P2P_INFO     _IOWR(TWCLML_IOCTL_MAGIC, 3, twclml_p2p_info_t)

/* 最大命令号 */
#define TWCLML_IOCTL_MAXNR 3

/* 函数声明 */
nvmlReturn_t nvmlInit(void);
nvmlReturn_t nvmlInit_v2(void);
nvmlReturn_t nvmlShutdown(void);
nvmlReturn_t nvmlDeviceGetCount(unsigned int* count);
nvmlReturn_t nvmlDeviceGetCount_v2(unsigned int* count);
nvmlReturn_t nvmlDeviceGetHandleByIndex(unsigned int index, nvmlDevice_t* device);
nvmlReturn_t nvmlDeviceGetPciInfo_v2(nvmlDevice_t device, nvmlPciInfo_t* pci);
nvmlReturn_t nvmlDeviceGetCudaComputeCapability(nvmlDevice_t device, int* major, int* minor);
nvmlReturn_t nvmlDeviceGetNvLinkState(nvmlDevice_t device, unsigned int link, nvmlEnableState_t* isActive);
nvmlReturn_t nvmlDeviceGetNvLinkRemotePciInfo(nvmlDevice_t device, unsigned int link, nvmlPciInfo_t* pci);
nvmlReturn_t nvmlDeviceGetP2PStatus(nvmlDevice_t device1, nvmlDevice_t device2, nvmlGpuP2PCapsIndex_t p2pIndex, nvmlGpuP2PStatus_t* p2pStatus);
const char* nvmlErrorString(nvmlReturn_t r);
nvmlReturn_t nvmlDeviceGetHandleByPciBusId(const char* pciBusId, nvmlDevice_t* device);
nvmlReturn_t nvmlDeviceGetIndex(nvmlDevice_t device, unsigned int* index);
nvmlReturn_t nvmlDeviceGetNvLinkCapability(nvmlDevice_t device, unsigned int link, nvmlNvLinkCapability_t capability, unsigned int* capResult);

#endif /* __TWCLML_DRV_H__ */